import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 加载本地 Iris 数据集
file_path = 'iris/iris.data'
columns = ['sepal_length', 'sepal_width', 'petal_length', 'petal_width', 'class']
data = pd.read_csv(file_path, header=None, names=columns)

# 提取特征（不包括类别）
X = data.iloc[:, :-1].values  # 取前四列作为特征


# K-Means 算法实现
def kmeans(X, k, max_iters=100):
    # 随机初始化中心点
    np.random.seed(42)
    centroids = X[np.random.choice(X.shape[0], k, replace=False)]

    for _ in range(max_iters):
        # 计算每个点到中心点的距离
        distances = np.linalg.norm(X[:, np.newaxis] - centroids, axis=2)

        # 分配每个点到最近的中心点
        labels = np.argmin(distances, axis=1)

        # 更新中心点
        new_centroids = np.array([X[labels == i].mean(axis=0) for i in range(k)])

        # 检查中心点是否收敛
        if np.all(centroids == new_centroids):
            break

        centroids = new_centroids

    return labels, centroids


# 设置簇的数量
k = 3
labels, centroids = kmeans(X, k)

# 可视化结果
plt.figure(figsize=(10, 6))
plt.scatter(X[:, 0], X[:, 1], c=labels, cmap='viridis', marker='o', edgecolor='k', s=100)
plt.scatter(centroids[:, 0], centroids[:, 1], c='red', marker='X', s=200, label='Centroids')
plt.title('K-Means Clustering on Iris Dataset')
plt.xlabel('Sepal Length')
plt.ylabel('Sepal Width')
plt.legend()
plt.grid()
plt.show()